import torch
from matplotlib import pyplot as plt

from config import Config as C
from model import resnet50
from utils import get_transform, getSingleImage


def main():
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    trans = get_transform()

    net = resnet50(num_classes=C.CLASSES_NUM)
    net.load_state_dict(torch.load(C.MODEL_WEIGHT_PATH, map_location='cpu'))
    net.to(device)
    net.eval()

    img_origin, label = getSingleImage(C.ROOT_PATH)
    img = trans['test'](img_origin)
    img = torch.unsqueeze(img, dim=0)
    with torch.no_grad():
        output = net(img.to(device))
        output = torch.max(output, dim=1)[1]
        true_idx = label.item()
        pred_idx = output.item()
        true_label = C.CLASSES[true_idx]
        pred_label = C.CLASSES[pred_idx]

    plt.imshow(img_origin)
    plt.xlabel(f"True Label:{true_label}   Predicate Label:{pred_label}")
    plt.title(f"The predicate is {'correct' if true_idx - pred_idx == 0 else 'incorrect'}")
    plt.show()


if __name__ == '__main__':
    main()
